﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ZGSharp.Optimeters
{
    public class OptBase
    {
        public virtual void opt(double[] paramter, double[] grad, double lr)
        {
        }
    }
    public class SGD : OptBase
    {
        public override void opt(double[] paramter, double[] grad, double lr)
        {
            Parallel.For(0, paramter.Length, delegate (int i)
            {
                paramter[i] = paramter[i] - lr * grad[i];
            });
        }
    }
    public class Momentum : OptBase
    {
        private Dictionary<double[], double[]> map_last_gard = new Dictionary<double[], double[]>();

        public override void opt(double[] paramter, double[] grad, double lr)
        {
            double[] last_gard = null;
            if (!map_last_gard.ContainsKey(paramter))
            {
                last_gard = new double[grad.Length];
                this.map_last_gard.Add(paramter, last_gard);
            }
            last_gard = this.map_last_gard[paramter];
            if (paramter.Length != last_gard.Length)
            {
                throw new ArgumentException("is error argv in Momentum");
            }
            Parallel.For(0, paramter.Length, delegate (int i)
            {
                double v = grad[i] * 0.9 + last_gard[i] * 0.1;
                paramter[i] = paramter[i] - lr * v;
                last_gard[i] = v;
            });
        }

    }
    public class AdaGrad : OptBase
    {
        private Dictionary<double[], double[]> map_last_gard = new Dictionary<double[], double[]>();
        public override void opt(double[] paramter, double[] grad, double lr)
        {
            double[] last_gard = null;
            if (!map_last_gard.ContainsKey(paramter))
            {
                last_gard = new double[grad.Length];
                this.map_last_gard.Add(paramter, last_gard);
            }
            last_gard = this.map_last_gard[paramter];
            if (paramter.Length != last_gard.Length)
            {
                throw new ArgumentException("is error argv in AdaGrad");
            }
            Parallel.For(0, paramter.Length, delegate (int i)
            {
                paramter[i] = paramter[i] - lr * grad[i] / Math.Sqrt(last_gard[i] + 0.0001);
            });
            Parallel.For(0, grad.Length, delegate (int i)
            {
                last_gard[i] += grad[i] * grad[i];
            });
        }

    }
    public class Adam : OptBase
    {
        private Dictionary<double[], double[]> map_last_gard = new Dictionary<double[], double[]>();
        public override void opt(double[] paramter, double[] grad, double lr)
        {
            double[] last_gard = null;
            if (!map_last_gard.ContainsKey(paramter))
            {
                last_gard = new double[grad.Length * 2];
                this.map_last_gard.Add(paramter, last_gard);
            }
            last_gard = this.map_last_gard[paramter];
            if (paramter.Length * 2 != last_gard.Length)
            {
                throw new ArgumentException("is error argv in Adam");
            }
            Parallel.For(0, grad.Length, delegate (int i)
            {
                double mt = 0.0;
                double vt = 0.0;
                double gt = grad[i];
                mt += 0.9 * last_gard[i] + 0.1 * gt;
                vt += 0.99 * last_gard[i + paramter.Length] + 0.01 * gt * gt;
                last_gard[i] = mt;
                last_gard[i + grad.Length] = vt;
            });
            Parallel.For(0, paramter.Length, delegate (int i)
            {
                paramter[i] = paramter[i] - lr * last_gard[i] / Math.Sqrt(last_gard[i + grad.Length] + 0.0001);
                if (double.IsNaN(paramter[i]))
                {
                    throw new ArgumentException("is NaN value in Adam");
                }
            });
        }
    }
}
